#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: pycharm
@file: prune_lenet.py
@time: 2020/10/15 11:53
@desc:https://zhuanlan.zhihu.com/p/129376087
'''

from torch_op.models.lenet import Lenet
import torch
import torch.nn.utils.prune as prune
from torchvision.models.resnet import resnet50
import torch.nn.functional as F
import torch.nn as nn
from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import DataLoader
from torch_op.common import get_sync_time

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Lenet()
weight = torch.load('torch_op/mnist/models_lenet/best.pt')['net']
model.load_state_dict(weight)
# model = resnet50(False,num_classes=10)
# model.load_state_dict(torch.load('torch_op/mnist/models/best.pt')['net'])

model = model.cuda()


test_dataset = MNIST('torch_op/mnist/data',train=False,download=True,
                      transform=transforms.Compose([
                          transforms.ToTensor(),
                          # transforms.Normalize((0.1307,),(0.3081))
                      ])
                      )

test_loader = DataLoader(test_dataset,batch_size=64,shuffle=False,num_workers=4)

def test():
    global best_acc
    model.eval()
    total_right = 0
    total_wrong = 0
    with torch.no_grad():
        for batch_idx,(data,target) in enumerate(test_loader):
            # if data.size(1) == 1:
            #     data = torch.cat([data, data, data], dim=1)
            data = data.cuda()
            target = target.cuda()
            output = model(data)
            output = output.max(dim=1)[1]
            right = torch.sum(output==target).item()
            wrong = len(data) - right
            total_right += right
            total_wrong += wrong
    acc = total_right/(total_wrong+total_right)
    print('acc:{}'.format(acc))

s = get_sync_time()
test()
e = get_sync_time()
print(e-s)

module = model.conv1
# print(list(module.named_modules()))
# print(list(module.named_buffers())) #mask
# print(list(module.named_parameters()))
#weights = weight_ori*mask

# prune.random_unstructured(module, name="weight", amount=0.0)
# prune.ln_structured(module,name='weight',amount=0.3,n=2,dim=0)
# print(list(module.named_buffers()))
# for hook in module._forward_pre_hooks.values():
#     if hook._tensor_name == "weight":
#        break
# print(list(hook))

# s = get_sync_time()
# test()
# e = get_sync_time()
# print(e-s)


# print(model.state_dict().keys())

# prune.remove(module, 'weight')
# print(list(module.named_parameters()))


#global pruning
model = Lenet()
weight = torch.load('torch_op/mnist/models_lenet/best.pt')['net']
model.load_state_dict(weight)
model.cuda()

parameters_to_prune = (
    (model.conv1, 'weight'),
    (model.conv2, 'weight'),
    (model.fc1, 'weight'),
    (model.fc2, 'weight'),
    (model.fc3, 'weight'),
)

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.9,
)

print(
    "Sparsity in conv1.weight: {:.2f}%".format(
        100. * float(torch.sum(model.conv1.weight == 0))
        / float(model.conv1.weight.nelement())
    )
)
print(
    "Sparsity in conv2.weight: {:.2f}%".format(
        100. * float(torch.sum(model.conv2.weight == 0))
        / float(model.conv2.weight.nelement())
    )
)
print(
    "Sparsity in fc1.weight: {:.2f}%".format(
        100. * float(torch.sum(model.fc1.weight == 0))
        / float(model.fc1.weight.nelement())
    )
)
print(
    "Sparsity in fc2.weight: {:.2f}%".format(
        100. * float(torch.sum(model.fc2.weight == 0))
        / float(model.fc2.weight.nelement())
    )
)
print(
    "Sparsity in fc3.weight: {:.2f}%".format(
        100. * float(torch.sum(model.fc3.weight == 0))
        / float(model.fc3.weight.nelement())
    )
)
print(
    "Global sparsity: {:.2f}%".format(
        100. * float(
            torch.sum(model.conv1.weight == 0)
            + torch.sum(model.conv2.weight == 0)
            + torch.sum(model.fc1.weight == 0)
            + torch.sum(model.fc2.weight == 0)
            + torch.sum(model.fc3.weight == 0)
        )
        / float(
            model.conv1.weight.nelement()
            + model.conv2.weight.nelement()
            + model.fc1.weight.nelement()
            + model.fc2.weight.nelement()
            + model.fc3.weight.nelement()
        )
    )
)

s = get_sync_time()
test()
e = get_sync_time()
print(e-s)